#!/usr/bin/env python3

import rospy
from sensor_msgs.msg import NavSatFix, NavSatStatus  # 确保导入了 NavSatStatus
from nav_msgs.msg import Odometry
from geometry_msgs.msg import Point, Pose, Quaternion, Twist, Vector3
from pyproj import Transformer, CRS
import math
import tf


class RTKOdometryNode:
    def __init__(self):
        rospy.init_node('rtk_odometry_node')

        # ROS参数
        self.frame_id = rospy.get_param('~frame_id', 'odom')
        self.child_frame_id = rospy.get_param('~child_frame_id', 'base_link')

        # 设置源（WGS84）和目标（UTM，广州位于UTM 50N）坐标系
        self.wgs84 = CRS("EPSG:4326")
        self.utm_zone = CRS("EPSG:32649")
        self.transformer = Transformer.from_crs(self.wgs84, self.utm_zone)

        # 上一个点的UTM坐标和时间
        self.previous_utm_x = None
        self.previous_utm_y = None
        self.previous_time = None

        # ROS发布者 (发布Odometry消息)
        self.odom_pub = rospy.Publisher('~odom', Odometry, queue_size=10)
        # tf 广播
        self.tf_broadcaster = tf.TransformBroadcaster()

        # ROS订阅者 (订阅NavSatFix消息)
        self.navsat_sub = rospy.Subscriber('rtk_reader_node/navsat', NavSatFix, self.navsat_callback)

    def navsat_callback(self, msg):
        """处理NavSatFix消息的回调函数"""
        # 检查是否是有效的RTK定位 (STATUS_FIX 或 STATUS_GBAS_FIX)
        if msg.status.status == NavSatStatus.STATUS_NO_FIX:
            rospy.logwarn("No RTK fix. Odometry not published.")
            return
        # 如果需要更高的精度, 可以只处理STATUS_GBAS_FIX
        # if msg.status.status != NavSatStatus.STATUS_GBAS_FIX:
        #    return

        current_time = msg.header.stamp

        # 转换为UTM坐标
        utm_x, utm_y = self.transformer.transform(msg.latitude, msg.longitude)

        # 计算速度和里程
        if self.previous_utm_x is not None and self.previous_utm_y is not None and self.previous_time is not None:
            delta_x = utm_x - self.previous_utm_x
            delta_y = utm_y - self.previous_utm_y
            delta_t = (current_time - self.previous_time).to_sec()

            if delta_t > 0:
                velocity_x = delta_x / delta_t
                velocity_y = delta_y / delta_t
                distance = math.sqrt(delta_x**2 + delta_y**2)
                yaw = math.atan2(delta_y, delta_x)

                # 发布Odometry消息
                self.publish_odometry(current_time, utm_x, utm_y, yaw, velocity_x, velocity_y)
                # 发布tf
                self.publish_tf(current_time, utm_x, utm_y, yaw)

        # 更新上一个点的数据
        self.previous_utm_x = utm_x
        self.previous_utm_y = utm_y
        self.previous_time = current_time

    def publish_odometry(self, time, x, y, yaw, vx, vy):
        """发布Odometry消息"""
        odom = Odometry()
        odom.header.stamp = time
        odom.header.frame_id = self.frame_id
        odom.child_frame_id = self.child_frame_id

        odom.pose.pose = Pose(Point(x, y, 0.0), self.create_quaternion_from_yaw(yaw))
        odom.twist.twist = Twist(Vector3(vx, vy, 0.0), Vector3(0.0, 0.0, 0.0))

        # 协方差 (根据实际情况设置)
        odom.pose.covariance = [0.01, 0, 0, 0, 0, 0,
                                 0, 0.01, 0, 0, 0, 0,
                                 0, 0, 0.01, 0, 0, 0,
                                 0, 0, 0, 99999, 0, 0,
                                 0, 0, 0, 0, 99999, 0,
                                 0, 0, 0, 0, 0, 99999]
        odom.twist.covariance = [0.01, 0, 0, 0, 0, 0,
                                  0, 0.01, 0, 0, 0, 0,
                                  0, 0, 0.01, 0, 0, 0,
                                  0, 0, 0, 99999, 0, 0,
                                  0, 0, 0, 0, 99999, 0,
                                  0, 0, 0, 0, 0, 99999]
        self.odom_pub.publish(odom)

    def publish_tf(self, time, x, y, yaw):
        """发布tf变换"""
        self.tf_broadcaster.sendTransform(
            (x, y, 0.0),
            self.create_quaternion_from_yaw(yaw),
            time,
            self.child_frame_id,
            self.frame_id
        )

    def create_quaternion_from_yaw(self, yaw):
        """从偏航角创建四元数"""
        q = tf.transformations.quaternion_from_euler(0, 0, yaw)
        return Quaternion(q[0], q[1], q[2], q[3])


if __name__ == '__main__':
    try:
        rtk_odometry_node = RTKOdometryNode()
        rospy.spin()
    except rospy.ROSInterruptException:
        pass


